﻿using System.Collections.Generic;
using UnityEngine;

namespace CWM.Skinn
{
    public static partial class SkinnEx
    {
        public static bool CloneOptimzeBones(this SkinnedMeshRenderer source, out SkinnedMeshRenderer cloned, bool log = true)
        {
            cloned = null;
            if (!source || !source.HasMinimumRequirements()) return false;

            Mesh mesh = source.sharedMesh.Clone() as Mesh;
            BoneWeight[] sourceWeights = mesh.boneWeights;
            Matrix4x4[] sourceBindPose = mesh.bindposes;
            Transform[] sourceBones = source.bones;
            List<Matrix4x4> bindposes = new List<Matrix4x4>();
            List<Transform> bones = new List<Transform>();

            if (source.rootBone && sourceBones.GetIndex(source.rootBone) > -1)
            {
                bones.Add(source.rootBone);
                bindposes.Add(sourceBindPose[sourceBones.GetIndex(source.rootBone)]);
            }

            for (int i = 0; i < sourceWeights.Length; i++)
            {
                BoneWeightData boneWeight = sourceWeights[i];
                boneWeight.Clamp();
                for (int ii = 0; ii < 4; ii++)
                {
                    float channelWeight = boneWeight.GetWeight(ii);
                    if (channelWeight > 0.0f)
                    {
                        int channelIndex = boneWeight.GetIndex(ii);
                        var bone = sourceBones[channelIndex];
                        if (!bone) { Debug.LogErrorFormat("Missing Bone :{0}", channelIndex); return false; }
                        int newIndex = bones.IndexOf(bone);
                        if (newIndex < 0)
                        {
                            newIndex = bones.Count;
                            bones.Add(bone);
                            bindposes.Add(sourceBindPose[channelIndex]);
                        }
                        boneWeight.SetIndex(ii, newIndex);
                    }
                    else boneWeight.SetIndex(ii, 0);
                }
                sourceWeights[i] = boneWeight;
            }

           
            cloned = source.Clone(true) as SkinnedMeshRenderer;

            if (!IsNullOrNotInAScene(source)) cloned.transform.SetParent(source.transform.parent);

            mesh.boneWeights = sourceWeights;
            mesh.bindposes = bindposes.ToArray();
            cloned.sharedMesh = mesh;
            cloned.bones = bones.ToArray();
            cloned.rootBone =  bones.FindMinHierarchicalOrder();
            cloned.SetBounds();

            if (log && Application.isEditor) Debug.LogFormat("{0}, original bone count: {1}. new bone count {2}", cloned, sourceBones.Length, cloned.bones.Length);
            return true;
        }


        public static bool CloneRemoveBones(SkinnedMeshRenderer source, Transform[] removedBones, out SkinnedMeshRenderer cloned)
        {
            cloned = null;
            if (!source || !source.HasMinimumRequirements() || source.bones.Length < 2) return false;

            var removedIndices = new List<int>(removedBones.Length);
            foreach (var bone in removedBones)
            {
                if (!bone) continue;
                var index = source.bones.GetIndex(bone.name); if (index < 0) continue;
                removedIndices.Add(index);
            }

            if (removedIndices.Count == 0) return false;

            var mesh = source.sharedMesh.Clone() as Mesh;
            var sourceWeights = mesh.boneWeights;
            var sourceBindposes = mesh.bindposes;
            var sourceBones = source.bones;

            var bones = new List<int>(source.bones.Length);
            for (int i = 0; i < sourceWeights.Length; i++)
            {
                BoneWeightData boneWeight = sourceWeights[i];
                boneWeight.Clamp();
                for (int ii = 0; ii < 4; ii++)
                {
                    var channelWeight = boneWeight.GetWeight(ii);
                    var channelIndex = boneWeight.GetIndex(ii);
                    if (removedIndices.Contains(channelIndex)) channelWeight = -1;
                    if (channelWeight > 0f)
                    {
                        int newIndex = bones.IndexOf(channelIndex);
                        if (newIndex < 0)
                        {
                            newIndex = bones.Count;
                            bones.Add(channelIndex);
                        }
                        boneWeight.SetIndex(ii, newIndex);
                        boneWeight.SetWeight(ii, channelWeight);
                    }
                    else
                    {
                        boneWeight.SetIndex(ii, 0);
                        boneWeight.SetWeight(ii, 0);
                    }
                }
                boneWeight.Clamp();

                sourceWeights[i] = boneWeight;
            }

            cloned = source.Clone(true) as SkinnedMeshRenderer;
            if (!SkinnEx.IsNullOrNotInAScene(source)) cloned.transform.SetParent(source.transform.parent);

            mesh.boneWeights = sourceWeights;
            cloned.sharedMesh = mesh;

            var newPoses = new Matrix4x4[bones.Count];
            var newBones = new Transform[bones.Count];
            for (int i = 0; i < newBones.Length; i++)
            {
                newBones[i] = sourceBones[bones[i]];
                newPoses[i] = sourceBindposes[bones[i]];
            }

            mesh.bindposes = newPoses;
            cloned.bones = newBones;
            cloned.rootBone = cloned.bones.FindMinHierarchicalOrder();
            cloned.SetBounds();

            if (Application.isEditor) Debug.LogFormat("{0}, original bone count: {1}. new bone count {2}", cloned, sourceBones.Length, cloned.bones.Length);
            return cloned;
        }


        public static bool CloneRemoveBones(SkinnedMeshRenderer source, Transform[] removedBones, Transform fallbackBone, out SkinnedMeshRenderer cloned)
        {
            cloned = null;
            if (!source || !source.HasMinimumRequirements() || source.bones.Length < 2) return false;

            var removedIndices = new List<int>(removedBones.Length);
            foreach (var bone in removedBones)
            {
                if (!bone) continue;
                var index = source.bones.GetIndex(bone); if (index < 0) continue;
                removedIndices.Add(index);
            }

            if (removedIndices.Count == 0) return false;

            var fallbackIndex = source.bones.GetIndex(fallbackBone); if (fallbackIndex < 0) return false;

            var mesh = source.sharedMesh.Clone() as Mesh;
            var sourceWeights = mesh.boneWeights;
            var sourceBindposes = mesh.bindposes;
            var sourceBones = source.bones;

            var bones = new List<int>(source.bones.Length);
            for (int i = 0; i < sourceWeights.Length; i++)
            {
                BoneWeightData boneWeight = sourceWeights[i];
                boneWeight.Clamp();

                var fallBackChannel = boneWeight.FindChannel(fallbackIndex);

                for (int ii = 0; ii < 4; ii++)
                {
                    var channelWeight = boneWeight.GetWeight(ii);
                    var channelIndex = boneWeight.GetIndex(ii);
                    if (removedIndices.Contains(channelIndex))
                    {
                        if (fallBackChannel < 0)
                        {
                            channelIndex = fallbackIndex;
                        }
                        else
                        {
                            int newIndex = bones.IndexOf(fallbackIndex);
                            if (newIndex < 0)
                            {
                                newIndex = bones.Count;
                                bones.Add(channelIndex);
                            }
                            boneWeight.SetIndex(ii, newIndex);
                            boneWeight.SetWeight(ii, boneWeight.GetWeight(fallBackChannel) + channelWeight);

                            continue;
                        }
                    }

                    if (channelWeight > 0f)
                    {
                        int newIndex = bones.IndexOf(channelIndex);
                        if (newIndex < 0)
                        {
                            newIndex = bones.Count;
                            bones.Add(channelIndex);
                        }
                        boneWeight.SetIndex(ii, newIndex);
                        boneWeight.SetWeight(ii, channelWeight);
                    }
                    else
                    {
                        boneWeight.SetIndex(ii, 0);
                        boneWeight.SetWeight(ii, 0);
                    }
                }
                boneWeight.Clamp();
                //boneWeight.Order();

                sourceWeights[i] = boneWeight;
            }

            cloned = source.Clone(true) as SkinnedMeshRenderer;
            if (!SkinnEx.IsNullOrNotInAScene(source)) cloned.transform.SetParent(source.transform.parent);

            mesh.boneWeights = sourceWeights;
            cloned.sharedMesh = mesh;

            var newPoses = new Matrix4x4[bones.Count];
            var newBones = new Transform[bones.Count];
            for (int i = 0; i < newBones.Length; i++)
            {
                newBones[i] = sourceBones[bones[i]];
                newPoses[i] = sourceBindposes[bones[i]];
            }

            mesh.bindposes = newPoses;
            cloned.bones = newBones;
            cloned.rootBone = cloned.bones.FindMinHierarchicalOrder();
            cloned.SetBounds();

            if (Application.isEditor) Debug.LogFormat("{0}, original bone count: {1}. new bone count {2}", cloned, sourceBones.Length, cloned.bones.Length);
            return cloned;
        }
    }
}
